#! -*- coding: utf-8 -*-
"""
@Author: AI
@Create Time: 20240625
@Info: 数据预处理
"""
import json
import random
from torch.utils.data import Dataset


class UIEDataSet(Dataset):
    """数据加载"""
    def __init__(self, file_path, tokenizer, is_train, extend_relation=False, max_length=512, neg_sampling=False):
        self.file_path = file_path
        self.tokenizer = tokenizer
        self.extend_relation = extend_relation
        self.is_train = is_train
        self.data = []
        self.neg_sampling = neg_sampling
        self.max_length = max_length

        self.read_json()

    def read_json(self):
        """读取文件"""
        with open(self.file_path, 'r', encoding='utf-8') as f:
            for line in f.readlines():
                line = json.loads(line.strip())
                text = line['text']
                if "entities" in line: # 判断该字段是否存在
                    label = line['entities']
                    for label_ in label:
                        label_.append(text[label_[0]:label_[1]])
                        label_[1] = label_[1] - 1
                    if label != []:
                        self.data.append(line)
        random.shuffle(self.data)

    @classmethod
    def find_pos(cls, text_tokens, start, end, entity):
        """找寻新的位置"""
        offset = -1
        start_matched, end_matched = False, False
        for i, token in enumerate(text_tokens):
            if token.startswith('##'):
                token = token.replace('##', '', 1)
            offset += len(token)

            if start == 0:
                start_matched = True

            if offset + 1 == start:
                start_matched = True
                start = i + 1

            if offset == end:
                end = i
                end_matched = True

        if not start_matched or not end_matched:
            print('don`t get position!!!')
            print(''.join(text_tokens), start, end, entity)

        return start, end

    def create_example(self, text, prompt, labels=None):
        """处理输入文本"""
        prompt_tokens = self.tokenizer.tokenize(prompt)
        text_tokens = self.tokenizer.tokenize(text)

        if len(prompt_tokens) + len(text_tokens) > self.max_length - 3:
            print(text)
            text_tokens = text_tokens[:self.max_length - 3 - len(prompt_tokens)]

        tokens = ['[CLS]'] + prompt_tokens + ['[SEP]'] + text_tokens + ['SEP']
        input_ids = self.tokenizer.convert_tokens_to_ids(tokens)
        token_type_ids = [0] * (len(prompt_tokens) + 2) + [1] * (len(text_tokens) + 1)
        mask_ids = [1] * len(input_ids)
        output = {
            'input_ids': input_ids,
            'token_type_ids': token_type_ids,
            'attention_mask': mask_ids,
            'seq_len': len(input_ids)
        }
        if self.is_train:
            start_label = [0] * len(input_ids)
            end_label = [0] * len(input_ids)
            if labels and isinstance(labels, list):
                for label in labels:
                    start, end, entity_type, entity = label
                    new_start, new_end = self.find_pos(text_tokens, start, end, entity)
                    start_label[len(prompt_tokens) + 2 + new_start] = 1
                    end_label[len(prompt_tokens) + 2 + new_end] = 1
            output.update({
                'start_label': start_label,
                'end_label': end_label
            })
        return output

    def __len__(self):
        return len(self.data)

    # def __getitem__(self, idx):
    #     line = self.data[idx]
    #     text = line['query']
    #     prompt = line['prompt']
    #     if self.is_train:
    #         label = line['span']
    #         inputs = self.create_example(text, prompt, label)
    #     else:
    #         inputs = self.create_example(text, prompt)
    #
    #     return inputs



    def __getitem__(self, idx):
        line = self.data[idx]
        text = line['text']
        if len(line['cats']) == 0:
            prompt = ''
        else:
            prompt = line['cats'][0]
        if self.is_train:
            label = line['entities']
            inputs = self.create_example(text, prompt, label)
        else:
            inputs = self.create_example(text, prompt)

        return inputs


